"""
Author: Xinyuan Ye
Description: Simple Backtracking Algorithm for Sudoku Solver Visualization.
             The program visually demonstrates the backtracking process to 
             solve a given Sudoku puzzle. Each step of the solving process 
             is animated with valid numbers shown in green and invalid numbers 
             (leading to backtracking) shown in red.

Usage:
1. Save the Sudoku puzzle as a text file (0 represents empty cells).
2. Run the program: python sudoku_visualization.py <sudoku_file.txt>
"""

import pygame
import time
import sys

# Define constants for the board
WIDTH, HEIGHT = 540, 540  # Window size
ROWS, COLS = 9, 9  # 9x9 Sudoku grid
CELL_SIZE = WIDTH // COLS  # Cell size
GRID_COLOR = (0, 0, 0)  # Black for grid lines
VALID_COLOR = (0, 255, 0)  # Green for valid number
INVALID_COLOR = (255, 0, 0)  # Red for invalid attempt
NUMBER_COLOR = (0, 0, 0)  # Default color for numbers

# Initialize pygame
pygame.init()

# Set the display size
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Sudoku Solver Visualization - CSP")

# Define the font
font = pygame.font.SysFont(None, 40)

# Read the Sudoku puzzle from a text file
def load_sudoku_from_file(filepath):
    with open(filepath, 'r') as f:
        puzzle = []
        for line in f:
            puzzle.append([int(num) for num in line.split()])
    return puzzle

# Draw the Sudoku board
def draw_board(board):
    # Draw the grid
    for i in range(ROWS + 1):
        thickness = 4 if i % 3 == 0 else 1  # Thicker lines for 3x3 boxes
        pygame.draw.line(screen, GRID_COLOR, (0, i * CELL_SIZE), (WIDTH, i * CELL_SIZE), thickness)
        pygame.draw.line(screen, GRID_COLOR, (i * CELL_SIZE, 0), (i * CELL_SIZE, HEIGHT), thickness)

    # Draw the numbers on the board
    for i in range(ROWS):
        for j in range(COLS):
            if board[i][j] != 0:
                text = font.render(str(board[i][j]), True, NUMBER_COLOR)
                screen.blit(text, (j * CELL_SIZE + 15, i * CELL_SIZE + 15))

# Improved: Check if placing num in board[row][col] is valid
def is_valid(board, num, row, col):
    # Check row
    for i in range(COLS):
        if board[row][i] == num:
            return False
    # Check column
    for i in range(ROWS):
        if board[i][col] == num:
            return False
    # Check 3x3 grid
    box_row = row // 3 * 3
    box_col = col // 3 * 3
    for i in range(3):
        for j in range(3):
            if board[box_row + i][box_col + j] == num:
                return False
    return True

# Solve the Sudoku board with backtracking
def solve_sudoku(board):
    for row in range(ROWS):
        for col in range(COLS):
            if board[row][col] == 0:  # Find an empty cell
                for num in range(1, 10):  # Try numbers 1-9
                    if is_valid(board, num, row, col):
                        board[row][col] = num
                        draw_current_cell(row, col, num, VALID_COLOR)  # Show the number in green for valid
                        pygame.display.update()
                        time.sleep(0.02)  # Optional delay for speed control
                        
                        if solve_sudoku(board):  # Recursively try to solve
                            return True
                        board[row][col] = 0  # Reset the cell if solution fails
                        draw_current_cell(row, col, num, INVALID_COLOR)  # Show the number in red for invalid
                        pygame.display.update()
                return False
    return True

# Highlight the current cell and show the number being tested
def draw_current_cell(row, col, num, color):
    # Draw cell background color
    pygame.draw.rect(screen, color, (col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE))
    # Draw the number being tested
    text = font.render(str(num), True, NUMBER_COLOR)
    screen.blit(text, (col * CELL_SIZE + 15, row * CELL_SIZE + 15))
    # Redraw grid lines to keep them visible
    pygame.draw.rect(screen, GRID_COLOR, (col * CELL_SIZE, row * CELL_SIZE, CELL_SIZE, CELL_SIZE), 1)

# Main visualization loop
def main():
    if len(sys.argv) < 2:
        print("Usage: python sudoku_visualization.py <sudoku_file.txt>")
        sys.exit(1)

    # Load the Sudoku puzzle from a text file
    filepath = sys.argv[1]
    sudoku_board = load_sudoku_from_file(filepath)

    running = True
    screen.fill((255, 255, 255))  # White background
    draw_board(sudoku_board)
    pygame.display.update()

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        solve_sudoku(sudoku_board)
        draw_board(sudoku_board)
        pygame.display.update()

    pygame.quit()

if __name__ == "__main__":
    main()